package com.it.you.common.aspect;

import com.google.common.util.concurrent.RateLimiter;
import com.it.you.common.annotation.RateLimit;
import com.it.you.common.annotation.RateLimited;
import com.it.you.common.utils.RateLimiterUtil;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @Author : YouHui
 * @create 2023/5/29 15:02
 */
@Aspect
@Component
@Order(1)
public class RateLimiterAspect {
    // 用于存储RateLimiter对象的缓存
    private final Map<String, RateLimiter> rateLimiterCache = new ConcurrentHashMap<>();

    @Before("@annotation(rateLimited)")
    public void limit(JoinPoint joinPoint, RateLimited rateLimited) throws Exception {
        // 获取注解中指定的最多处理的请求数量和时间间隔
        int permits = rateLimited.value();
        int duration = rateLimited.duration();

        // 从缓存中获取对应的RateLimiter对象
        String key = joinPoint.getSignature().toString() + "_" + duration;
        RateLimiter limiter = rateLimiterCache.computeIfAbsent(key, k -> RateLimiter.create(permits, duration, TimeUnit.SECONDS));

        // 尝试获取令牌
        if (!limiter.tryAcquire(permits,duration,TimeUnit.SECONDS)) {
            // 请求被限流，抛出异常
            throw new Exception("Too many requests");
        }
    }

    @Around("@annotation(rateLimit)")
    public Object limit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        String key = joinPoint.getSignature().toString();
        RateLimiter rateLimiter = RateLimiterUtil.getRateLimiter(key);
        if (rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.limit(), TimeUnit.SECONDS)) {
            return joinPoint.proceed();
        } else {
            return "Limit!";
        }
    }

}
